// Copyright (c) MLLM Team.
// Licensed under the MIT License.

#pragma once

#include "mllm/backends/qnn/op/QNNBaseOp.hpp"
#include "mllm/compile/ir/linalg/Op.hpp"
#include "mllm/compile/ir/Node.hpp"
#include "mllm/compile/ir/tensor/Value.hpp"
#include "mllm/core/aops/X2XOp.hpp"
#include <vector>

namespace mllm::qnn {

class QNNX2XOp final : public aops::X2XOp {
 public:
  explicit QNNX2XOp(const aops::X2XOpOptions& options);

  void forward(const std::vector<Tensor>& inputs, std::vector<Tensor>& outputs) override;
};

class QNNX2XOpFactory : public TypedOpFactory<OpTypes::kX2X, aops::X2XOpOptions> {
 public:
  std::shared_ptr<BaseOp> createOpImpl(const aops::X2XOpOptions& options) override {
    return std::make_shared<QNNX2XOp>(options);
  }
};

// A placeholder, it should not occur in QNN graph
class QNNX2XPattern : public QNNBasePattern {
 public:
  bool isMatch(const ir::op_ptr_t& op) override { return op->isa_<ir::linalg::X2XOp>(); }

  bool addNode(const std::string& graphName, const ir::op_ptr_t& op, const std::vector<ir::tensor::TensorValue::ptr_t>& inputs,
               const std::vector<ir::tensor::TensorValue::ptr_t>& outputs) override;

  static std::pair<OpTypes, std::shared_ptr<QNNX2XPattern>> create() {
    return {OpTypes::kX2X, std::make_shared<QNNX2XPattern>()};
  }
};

}  // namespace mllm::qnn